# PyGameNet, VERSION 0.3
# By Zach Barth
# Altered by Noah Kantrowitz
# September 28, 2007
import array
import socket
import threading
import time
import itertools
import struct
import errno
from collections import deque
from pickle import *
from select import select

try:
    import pygame
except ImportError:
    pygame = None
    
# Protocol type constants
PROTOCOL_SET_ID = 0x01
PROTOCOL_ADD_MESSAGE = 0x02
PROTOCOL_ADD_PLAYER = 0x03
PROTOCOL_SET_PLAYER_FIELD = 0x04
PROTOCOL_DEL_PLAYER = 0x05
PROTOCOL_SET_GAME_ARRAY = 0x06
PROTOCOL_SET_GAME_FIELD = 0x07
PROTOCOL_SEND_MESSAGE = 0x08
PROTOCOL_SEND_BROADCAST = 0x09

_protocol_names = dict((v, k[9:]) for k, v in vars().iteritems() if k.startswith('PROTOCOL_'))


class PyGameNetError(Exception):
    """Top-level exception for for PyGameNet."""
    
class NotConnectedError(PyGameNetError):
    """Error due to lack of connection."""
    
class AlreadyConnectedError(PyGameNetError):
    """Error due to already being connected."""
    
class ConnectionFailed(PyGameNetError):
    """Error due to attempted connection failing."""
    
class NotEnoughDataError(PyGameNetError):
    """Error due to not enough data being recieved."""
    
class ProtocolError(PyGameNetError):
    """Error in protocol data stream."""
    
class SocketEOF(Exception):
    """A marker meaning a SocketBuffer reached EOF."""

class Clock(object):
    """A replacement for pygame.time.Clock that allow for other threads."""
    
    def __init__(self):
        self._clock = pygame.time._Clock()
    def tick(self, rate):
        ret = self._clock.tick(rate)
        pygame.time.wait(2)
        return ret

    def get_fps(self):
        return self._clock.get_fps()
        
if pygame:
    pygame.time._Clock = pygame.time.Clock
    pygame.time.Clock = Clock

class NetBuffer(list):
    """Used to store data to be sent out over and pulled in from the network."""
    
    def __init__(self, initial_value=''):
        super(NetBuffer, self).__init__(initial_value)
        
    @classmethod
    def from_args(cls, proto, args):
        """Build a buffer from a given protocol type and data."""
        self = cls(chr(proto))
        self.push_short(0) # Placeholder for the length bytes
        for arg in args:
            self.push(arg)
            
        length = len(self) - 3 # Don't count the 1-byte proto and 2-byte len
        self[1] = chr(length // 256)
        self[2] = chr(length % 256)
        return self
        
    @classmethod
    def from_message(cls, client, msg, proto_type=PROTOCOL_SEND_MESSAGE):
        """Build a buffer for a given message."""
        msg_id, msg_types = client._message_types[msg.type]
        self = cls(chr(proto_type))
        self.push_short(0) # Placeholder for the length bytes
        
        # Add the message ID
        self.push_short(msg_id)
        
        # Add our arguments
        for code, arg in itertools.izip(msg_types, msg):
            self.type_codes[code][0](self, arg)
#        try:
#            if msg_types != '':
#                self.append(struct.pack(msg_types, *msg))
#        except struct.error, e:
#            print "Message:", msg
#            raise e

        # Finalize and return
        length = len(self) - 3 # Don't count the 1-byte proto and 2-byte len
        self[1] = chr(length // 256)
        self[2] = chr(length % 256)
        return self
            
    def push_string(self, data):
        """Add a string, preceeded by its length."""
        self.push_short(len(data))
        self.extend(data)

    def push_short(self, data):
        """Add a signed short, MSB first."""
        #I don't want to have to use pack and unpack, so for now, MSB will simply have 128 added to it. Thus, 128=0, 131=3, 127=-1.
        self.append(chr(data // 256 + 128))
        self.append(chr(data % 256))
        
    def push_float(self, data):
        """Add a float."""
        self.extend(struct.pack('d', data))
        
    def push(self, data):
        """A dispatching convenience function."""
        if isinstance(data, str):
            self.push_string(data)
        elif isinstance(data, int):
            self.push_short(data)
        elif isinstance(data, float):
            self.push_float(data)
        else:
            raise PyGameNetError('Invalid data for sending %r'%data)
        
    def pop_string(self):
        length = self.pop_short()
        ret = self[:length]
        del self[:length]
        return ''.join(ret)
        
    def pop_short(self):
        msb = ord(self.pop(0))
        lsb = ord(self.pop(0))
        return ((msb - 128) << 8) + lsb
        
    def pop_float(self):
        ret = struct.unpack('d', ''.join(self[:8]))[0]
        del self[:8]
        return ret
        
    def to_string(self):
        """Flatten to a string."""
        return ''.join(self)
        
    type_codes = {
        'h': (push_short, pop_short),
        's': (push_string, pop_string),
        'f': (push_float, pop_float),
    }
    
    def pop_by_code(self, code):
        """Pop an argument out of the buffer based on typecode."""
        return self.type_codes[code][1](self)
        
class SocketBuffer(object):
    """An enhanced socket wrapper that can deal with more complex buffereing."""

    def __init__(self, sock):
        self.sock = sock
        self.sock.setblocking(False)
        self.buf = deque()
        
    def get(self, n_bytes):
        """Return an iterable of the requested bytes."""
        while len(self.buf) < n_bytes:
            # As long as we don't have enough data ready
            try:
                data = self.sock.recv(4096)
                if data == '':
                    # EOF
                    raise SocketEOF
                self.buf.extend(data)
            except socket.error, e:
                if e[0] == errno.EWOULDBLOCK:
                    # No more data, must not have enough
                    return None
                else:
                    raise
            
        ret = [] 
        for i in xrange(n_bytes):
            # Yield the wanted data
            ret.append(self.buf.popleft())
        return ret
    
    def send(self, bytes):
        """Send all bytes, blocking if needed."""
        n_bytes = len(bytes)
        i = 0
        while i < n_bytes:
            try:
                i += self.sock.send(bytes[i:])
            except socket.error, e:
                if e[0] == errno.EWOULDBLOCK:
                    pass

    # For compat
    sendall = send

class Message(list):
    """A model for a single message on the network."""
    
    def __init__(self, type=None, id=None):
        super(Message, self).__init__()
        self.type = type
        self.id = id
        
    def __repr__(self):
        return '<pygamenet.Message type=%s %s %r>'%(self.type, self.id is not None and 'id=%s'%self.id or '', list(self))


class Client(object):
    """Network client."""
    
    def __init__(self):
        self.players = {}
        self.game = {}
        self.id = None
        self._message_types = {} # name -> (id, type_codes)
        self._message_types_inv = {} # id -> (name, type_codes)
        #self._thread = _ClientThread(self)
        #self._message_queue = Queue.Queue()
        self.sock = None
        self.buf = None
        #self.timeout = 30
        self._msg_iter = self._get()
        self._ready = False
        
    connected = property(lambda self: self.sock is not None)
    
    def send(self, msg):
        """Send a message to the server."""
        if self.sock is None:
            return
        if msg.type not in self._message_types:
            raise ValueError('Unknown message type %s'%msg.type)
            
        self.debug('Sending message %r', msg)
        buf = NetBuffer.from_message(self, msg)
        self.sock.send(buf.to_string())
        
    def broadcast(self, msg):
        """Send a message to all other clients and the server."""
        if self.sock is None:
            return
        if msg.type not in self._message_types:
            raise ValueError('Unknown message type %s' % msg.type)
            
        buf = NetBuffer.from_message(self, msg, proto_type=PROTOCOL_SEND_BROADCAST)
        self.sock.send(buf.to_string())
    
    def _get(self):
        """An iterator for looking for messages. Internal use only."""
        while True:
            try:
                # Look for a message
                while 1:
                    bytes = self.buf.get(3)
                    if bytes is not None:
                        break
                    yield None
                
                # Decode the metadata
                proto_type, proto_len_msb, proto_len_lsb = bytes
                proto_len = ord(proto_len_msb) * 256 + ord(proto_len_lsb)
                self.debug('Looking for %r bytes', proto_len)
            
                # Get the rest of the message
                while 1:
                    bytes = self.buf.get(proto_len)
                    if bytes is not None:
                        break
                    yield None
            
            except SocketEOF:
                # Do something
                raise Exception('Hmm')
            
            # Load up the message
            proto = NetBuffer(bytes)
            
            # Decode the message
            proto_type = ord(proto_type)
            self.debug('Got protocol type %s', _protocol_names.get(proto_type, '[unknown]'))
            if proto_type == PROTOCOL_SEND_MESSAGE:
                # A normal message. Further decode to a Message and push to queue.
                msg_type = proto.pop_short()
                if msg_type == 0:
                    # Mark that we are set to go
                    self._ready = True
                elif not self._ready:
                    # Ignore messages before we get an ack from the server
                    continue
                
                if msg_type not in self._message_types_inv:
                    raise ProtocolError('Unknown message type %s'%msg_type)
                msg_type_name, msg_type_codes = self._message_types_inv[msg_type]
                self.debug('Got message type %s<%s>', msg_type_name, msg_type)
                msg = Message(msg_type_name)
                for c in msg_type_codes:
                    msg.append(proto.pop_by_code(c))
                yield msg
            elif proto_type == PROTOCOL_SET_ID:
                self.id = proto.pop_short()
                self.debug('I am client %s', self.id)
            elif proto_type == PROTOCOL_ADD_MESSAGE:
                self.debug('Adding new message type')
                type_name = proto.pop_string()
                type_id = proto.pop_short()
                type_codes = proto.pop_string()
                self.debug('Got name=%r, id=%r, types=%r', type_name, type_id, type_codes)
                self._message_types[type_name] = (type_id, type_codes)
                self._message_types_inv[type_id] = (type_name, type_codes)
            elif proto_type == PROTOCOL_ADD_PLAYER:
                player_id = proto.pop_short()
                self.players[player_id] = {}
                # XXX: Communicate this to the Client <NPK>
            elif proto_type == PROTOCOL_DEL_PLAYER:
                 player_id = proto.pop_short()
                 del self.players[player_id]
                 # XXX: Communicate this to the Client <NPK>
            else:
                raise ProtocolError('Unknown protocol type %s'%proto_type)
            
    def poll(self):
        """Returns the next message in the message queue, or None if there are no messages."""
        if not self.connected:
            return None
        return self._msg_iter.next()
    
    def get(self):
        """Returns an iterable of all messages from the message queue."""
        msg = self.poll()
        while msg is not None:
            yield msg
            msg = self.poll()
    
    def set_player_field(self, field, value):
        """Request to set a player field for self; the server gets it in the form of a "client_set_player_field" message."""
        self.send("client_set_player_field", (field, value))        
    
    def set_game_field(self, field, value):
        """Request to set a game field for self; the server gets it in the form of a "client_set_game_field" message."""
        self.send("client_set_game_field", (field, value))    
    
    def connect(self, ip, port):
        """Connect to a server at ip:port."""
        if self.sock is not None:
            raise AlreadyConnectedError
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            self.sock.connect((ip, port))
        except socket.error, e:
            if e[0] == errno.ECONNREFUSED:
                self.sock = None
                raise ConnectionFailed
            raise
        self.sock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
        self.buf = SocketBuffer(self.sock)
        #self._thread.start()
    
    def disconnect(self):
        """Disconnect from a server."""
        if self.sock is None:
            return
        self.sock.shutdown(socket.SHUT_RDWR)
        #self._thread.stop = True
        #self._thread.join()
        #self._thread = _ClientThread(self)
        self.sock = None
        self.buf = None 
        self.id = None
        
    def debug(self, msg, *args):
        """Debugging output hook."""
        # XXX: Use the logging module. <NPK>
        print ('PyGameNet: '+msg)%args
        pass
    
    def wait_ready(self):
        """Block until the recipt of a server_ready message."""
        while 1:
            for msg in self.get():
                if msg.type == 'server_ready':
                    return


class Server(object):
    """A network server."""
    
    def __init__(self, port, message_types=()):
        """Create a Server, listening on port for connections."""
        self.port = port
        self.players = {}
        self.game = {}
        # self._threads = {}
        # self._message_queue = Queue.Queue()
        # self._master_thread = _ServerMasterThread(self)
        self.timeout = 30
        
        message_types = itertools.chain((('server_ready', ''), ('client_set_player_field', ''), ('client_set_game_field', ''), ('client_broadcast', '')), message_types)
        self._message_types = {}
        self._message_types_inv = {}
        for (name, codes), id in itertools.izip(message_types, itertools.count()):
            self._message_types[name] = (id, codes)
            self._message_types_inv[id] = (name, codes)
        
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.bind(('', self.port))
        self.sock.listen(5)
        self.sock.setblocking(False)
        
        self._master_task = self._master_task()
        self._tasks = deque()
        
    def _broadcast(self, data):
        """Send data to all client sockets."""
        for id, player in self.players.items():
            try:
                player['buf'].sendall(data)
            except socket.error, e:
                if e[0] == errno.EPIPE:
                    # Broken pipe, connection lost
                    self.debug('Connection to client %s lost', id)
                    del self.players[id]
                    continue
                raise
    
    def send(self, id, msg):
        """Send message to player by id."""
        buf = NetBuffer.from_message(self, msg)
        self.players[id]['buf'].send(buf.to_string())
    
    def broadcast(self, msg):
        """Send message to all players."""
        buf = NetBuffer.from_message(self, msg)
        self._broadcast(buf.to_string())
    
    def poll(self):
        """Returns the next message in the message queue, or None if there are no messages."""
        # Pump waiting connections
        self._master_task.next()
        for i in xrange(len(self._tasks)):
            task, id = self._tasks.popleft()
            if id in self.players:
                self._tasks.append((task, id))
                msg = task.next()
                if msg is not None:
                    return msg
        return None
    
    def get(self):
        """Returns an iterable of all messages from the message queue."""
        msg = self.poll()
        while msg is not None:
            yield msg
            msg = self.poll()
    
    # def set_player_field(self, id, field, value):
    #     """Set a globally synchronized player field for player by id (sends to all clients)."""
    #     self._players[id][field] = value
    #     #Send this to the Clients.
    #     buf = NetBuffer()
    #     buf.push_int(PROTOCOL_SET_PLAYER_FIELD)
    #     buf.push_int(id)
    #     buf.push_string(dumps(field, HIGHEST_PROTOCOL))
    #     buf.push_string(dumps(value, HIGHEST_PROTOCOL))
    #     self._super_thread.protocol_broadcast(buf.to_string())
    # 
    # def set_game_field(self, field, value):
    #     """Set a globally synchronized game field."""
    #     self._game[field] = value
    #     #Send this to the Clients.
    #     buf = NetBuffer()
    #     buf.push_int(PROTOCOL_SET_GAME_FIELD)
    #     buf.push_string(dumps(field, HIGHEST_PROTOCOL))
    #     buf.push_string(dumps(value, HIGHEST_PROTOCOL))
    #     self._super_thread.protocol_broadcast(buf.to_string())
    # 
    # def disconnect_client(self, id):
    #     """Disconnect a client by ID."""
    #     self._threads[id].disconnect()
    
    def _get_ordinal(self, type):
        """Return the ordinal (numerical index) of the type string. Used for building network packets."""
        for i in range(0, len(self._message_types)):
            if self._message_types[i] == type:
                return i
        raise NameError, "Message type not found: %s" % type
    
    def get_server_info(self):
        """Get information about the server, such as IP address and host name."""
        info = socket.gethostbyaddr(socket.gethostname())
        return (info[2][0], info[0])
    
    def shutdown(self):
        """Shuts down all threads to allow program to properly close."""
#        self._master_thread.stop = True
        self.sock.shutdown(socket.SHUT_RDWR)
#        for t in self._threads.itervalues():
#            t.stop = True
#            t.sock.shutdown(socket.SHUT_RDWR)
#        socket.socket.shutdown(socket.SHUT_RDWR)
    
    def debug(self, msg, *args):
        pass
        print msg%args
    
    def _master_task(self):
        run = True
        while run:
            # Get an incoming connection
            try:
                sock, remote = self.sock.accept()
            except socket.error, e:
                if e[0] == errno.EWOULDBLOCK:
                    yield None
                    continue
                raise
                
            # Set connection parameters
            sock.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
            
            # Spawn worker task
            for new_id in itertools.count():
                if new_id not in self.players:
                    break
            task = self._worker_task(sock, remote, new_id)
            self.players[new_id] = {
                'sock': sock,
                'task': task,
                'buf': SocketBuffer(sock),
            }
            self._tasks.append((task, new_id))
            
    def _worker_task(self, sock, remote, client_id):
        sockbuf = SocketBuffer(sock)
        run = True
        
        # Inform the Server of our presence.
        msg = Message('client_connect', client_id)
        msg.id = client_id
        msg.append(remote)
        yield msg

        data = []

        # Push all message types to the new client
        for name, (id, codes) in self._message_types.iteritems():
            buf = NetBuffer.from_args(PROTOCOL_ADD_MESSAGE, (name, id, codes))
            data.extend(buf) # Send these all together at the same time

        # Send the client its ID
        buf = NetBuffer.from_args(PROTOCOL_SET_ID, (client_id,))
        data.extend(buf)

        # Send the client all other players
        for client in self.players:
            if client != client_id:
                buf = NetBuffer.from_args(PROTOCOL_ADD_PLAYER, (client,))
                data.extend(buf)

        # Transmit!
        sockbuf.sendall(''.join(data))

        #Inform the other Clients of our presence.
        buf = NetBuffer.from_args(PROTOCOL_ADD_PLAYER, (client_id,))
        self._broadcast(buf.to_string())

        #Let the Client know we're ready.
        self.send(client_id, Message('server_ready'))
        
        while run:
            try:
                # Block forever waiting for a start of message
                self.debug('Looking for a message')
                while 1:
                    bytes = sockbuf.get(3)
                    if bytes is not None:
                        break
                    yield None
                proto_type, proto_len_msb, proto_len_lsb = bytes
            
                # Get the rest of the message
                proto_len = ord(proto_len_msb) * 256 + ord(proto_len_lsb)
                self.debug('Waiting for %s bytes', proto_len)
                while 1:
                    bytes = sockbuf.get(proto_len)
                    if bytes is not None:
                        break
                    yield None
                
                proto = NetBuffer(bytes)
            except SocketEOF:
                run = False
                break
            except socket.error as (errno, strerror):
                self.debug("Connection dropped due socket.error:({0}): {1}".format(errno, strerror))
                run = False
                break

            # Decode the message
            proto_type = ord(proto_type)
            self.debug('Got protocol type %s', _protocol_names.get(proto_type, '[unknown]'))
            if proto_type == PROTOCOL_SEND_MESSAGE:
                # A normal message. Further decode to a Message and push to queue.
                msg_type = proto.pop_short()
                if msg_type not in self._message_types_inv:
                    raise ProtocolError('Unknown message type %s'%msg_type)
                msg_type_name, msg_type_codes = self._message_types_inv[msg_type]
                self.debug('Got message type %s<%s>', msg_type_name, msg_type)
                
                msg = Message(msg_type_name)
                msg.id = client_id
                for c in msg_type_codes:
                    msg.append(proto.pop_by_code(c))
                yield msg
            elif proto_type == PROTOCOL_SEND_BROADCAST:
                # A client multipoint broadcast
                # Reform and send to all other client
                len_msb = chr(len(proto) // 256)
                len_lsb = chr(len(proto) % 256)
                proto[0:0] = [chr(PROTOCOL_SEND_MESSAGE), len_msb, len_lsb]
                data = proto.to_string()
                for id, player in self.players.items():
                    if id != client_id:
                        try:
                            player['buf'].sendall(data)
                        except socket.error, e:
                            if e[0] == errno.EPIPE:
                                self.debug('Connection to client %s lost', client)
                                del self.players[id]
                                continue
                            raise
                                
                del proto[:3]
                
                # Process as normal and ship off to the server
                msg_type = proto.pop_short()
                if msg_type not in self._message_types_inv:
                    raise ProtocolError('Unknown message type %s'%msg_type)
                msg_type_name, msg_type_codes = self._message_types_inv[msg_type]
                msg = Message(msg_type_name)
                msg.id = client_id
                for c in msg_type_codes:
                    msg.append(proto.pop_by_code(c))
                yield msg
            else:
                raise ProtocolError('Unknown protocol type %s'%proto_type)
                
        del self.players[client_id]


if __name__ == '__main__':
    raise Exception
    # s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # s.connect(('127.0.0.1', 9124))
    # s.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
    # b = SocketBuffer(s)
    # print list(b.get(4, 5))
    # b = NetBuffer()
    # 
    # # Send P_SET_ID, 0x00, 0x02, 1
    # b.extend(chr(PROTOCOL_SET_ID)+'\x00\x02')
    # b.push(1)
    # 
    # # Send P_ADD_MESSAGE, 0x00, 0x09, 'foo', 1, 's'
    # b.extend(chr(PROTOCOL_ADD_MESSAGE)+'\x00\x0A')
    # b.push('foo')
    # b.push(1)
    # b.push('s')
    # 
    # # Send P_SEND_MESSAGE, 0x00, 0x09, 1, 'hello'
    # b.extend(chr(PROTOCOL_SEND_MESSAGE)+'\x00\x09')
    # b.push(1)
    # b.push('hello')
    # 
    # s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # s.bind(('127.0.0.1', 9123))
    # s.listen(5)
    # c, _ = s.accept()
    # c.send(b.to_string())
    s = Server(9124, [('die', '')])
    run = True
    while run:
        for msg in s.get_all():
            if msg.type == 'die':
                print 'Die!'
                run = False
    s.shutdown()